import bisect
import math
import string

from matplotlib import pyplot
import numpy
from scipy import optimize

from fitstat import FitStatistics

class FitException(Exception):
    '''
    An exception object specifically for fitting.
    '''
    pass

class Fitter(object):
    """
    An object to wrap the scipy fitting routines into a more friendly interface.
    ``Fitter`` is a base class that defines common code and an interface for fitters
    implementing specific functions.

    
    
    Public instance variables:
    --------------------------

    x             -- X-coordinates of the data to be fit
    y             -- Y-coordinates of the data to be fit
    weights       -- The reciprocal of the error in Y array
    stats          -- ``FitStatistics`` object containing the results of the last fit
    par_names     -- A list containing the names of the fit parameters
    npar          -- The number of parameters used in the fit
    fixedp        -- A list of indices of parameters that are not allowed to vary during a fit
    fixed_bg_pars -- A list of parameters values to supply to a fixed background function, if any.
    data_len      -- The number of points to be fit
    title         -- A string with title to be printed on the drawn fit
    xlabel        -- A string with the label for the x-axis
    ylabel        -- A string with the label for the y-axis

    """

    def __init__(self, x, y, par_names, yerr=None, xerr=None, fcn_str='1', title='My Fit', xlabel='', ylabel=''):
        """
        The default constructor for a ``Fitter`` object. A user would probably never call this
        directly--instead he or she would use one of the derived classes implementing a specific
        fit function.

        Arguments:

        x         -- X-coordinates of the data to be fit
        y         -- Y-coordinates of the data to be fit
        xerr      -- An array of errors in the X measurement
        yerr      -- An array of errors in the Y measurement
        par_names -- List of names for the parameters
        fcn_str   -- A Python-compilable string specifying the fit function
        """

        if len(x) < 3:
            raise FitException("Length-{0} X array supplied to Fitter constructor.".format(len(x)))
        if len(y) < 3:
            raise FitException("Length-{0} Y array supplied to Fitter constructor.".format(len(y)))
        if len(x) != len(y):
            raise FitException("Attempt to initialize fitter with unequal-length arrays (nx = {0}, ny = {1})".format(len(x), len(y)))
        if xerr is not None:
            raise NotImplementedError("I dunno how to do a fit with errors in X.")
        if yerr is not None:
            if len(yerr) != len(y):
                raise FitException("Attempt to initialize fitter with unequal-length arrays (yerr = {0}, y = {1})".format(len(yerr), len(y)))
            
            for i in xrange(len(yerr)):
                if math.fabs(yerr[i]) < 1e-8:
                    raise FitException("Attempt to initialize fitter with errors on at least one measurement equal to zero (element {0})".format(yerr[i]))

        
            self.weights = 1.0 / numpy.array(yerr)
        else:
            self.weights = numpy.ones((len(y),))

        self.stats             = None
        '''``FitStatistics`` object holding results of the last fit'''
        self.par_names        = par_names
        '''A list of parameter names for the fit function in ``fcn_str``'''
        self._par_text_offset = 0.375
        '''An offset used in positioning the text of the fitted parameters on a plot'''
        self.npar             = len(self.par_names)
        '''Number of parameters in the fit.'''
        self._fmt             = string.Formatter()
        '''Used to format ``fcn_str`` whenever it is compiled'''

        self._reset_fcn(fcn_str)
        
        self.fixedp = list()
        '''A list of the parameters that must remained fixed for the next fit'''
        self.x = x
        '''The X-coordinates of the data to be fit'''
        self.y = y
        '''The Y-coordinates of the data to be fit'''
        self.ext = 'png'
        '''The extension of the image file created when a fit is drawn.'''

        self.fixed_bg_pars = None
        '''Parameters for use with a fixed background function, if present'''
        self.title = title
        '''Title for drawing the fit'''
        self.xlabel = xlabel
        '''X-axis label for drawing the fit'''
        self.ylabel = ylabel
        '''Y-axis label for drawing the fit'''

        
    def _reset_fcn(self, fcn_str):
        """Set the current fit function to the compiled value of ``fcn_str``."""
        self._curr_fcn_str = "lambda p, x: " + self._fmt.vformat(fcn_str, dict([[i, 'p[{0}]'.format(i)] for i in xrange(self.npar)]), None)
        self.fcn = eval(self._curr_fcn_str)

    def fcn(self, p, x):
        '''
        The fit function without any fixed background contribution.

        Arguments:
        
        p -- parameters for the fit function
        x -- a single point or an array to be evaluated with the given parameters
        '''
        raise NotImplementedError("Fitter.fcn() was not set automatically when the fitter is constructed.")
    
    def _fcn(self, p, x):
        '''
        The fit function including the fixed background contribution, if supplied. This
        is the function called from ``Fitter.chi()``

        Arguments:
        
        p -- parameters for the fit function
        x -- a single point or an array to be evaluated with the given parameters
        '''
        return self.fcn(p, x) + self.fixed_bg(self.fixed_bg_pars, x)

    def fixed_bg(self, p, x):
        '''
        The function representing a fixed parameterized background. If no background function has
        been set, this function returns 0.

        Arguments:

        p -- parameters to the background function
        x -- a single point or an array to be evaluated with the given parameters
        '''
        
        return 0.0;

    def add_fixed_bg(self, fcn, pars):
        '''
        Offset the fit function during the fit by a background function ``fcn`` with fixed parameters ``pars``.

        Arguments:

        fcn  -- a Python function object returning the background when supplied a parameter list and x-coordinates (either a single point or list).
            The function should take two non-keyword arguments: first, the parameter list, and second, the x-coordinate or x-coordinate array.
        pars -- the fixed parameters at which this function should be evaluated during the fit.
        '''
        
        self.fixed_bg = fcn
        self.fixed_bg_pars = pars

    def clear_fixed_bg(self):
        '''
        Remove a fixed background function.

        '''
        
        self.fixed_bg = lambda p, x: 0.0
        self.fixed_bg_pars = None

    def fix_parameter(self, p, value=None):
        '''
        Mark a parameter as fixed for subsequent fits. 

        Arguments:

        p      -- p may be an integer index to the parameter to be fixed, in which case the second argument must be supplied,
            or a dictionary with indices of parameters to be fixed as keys and the values they are to be fixed to as the values of the dictionary.
        value  -- the value the parameter p should be fixed to, if p is a single parameter index
        '''
        
        if value is None:
            if type(p) != type(dict()):
                raise ArgumentError("BinnedFitter.fix_parameter() expects a dictionary when passed a single argument.")
            self._fix_p(p)
        else:
            self._fix_p({p: value})
        
    def _fix_p(self, pdict):
        '''
        Function for recompiling the function string when a parameter is fixed.
        '''
        
        vfmt_args = dict()
        self.fixedp = dict(pdict.items())
        for i in xrange(npar):
            if i in self.fixedp:
                vfmt_args[i] = str(self.fixedp[i])
            else:
                vfmt_args[i] = 'p[{0}]'.format(i)
        
        self._curr_fcn_str = self._fmt.vformat(self.fcn_str, vfmt_args)
        self.fcn = eval(self._curr_fcn_str)

    def release_parameter(self, p=None):
        '''
        Release a previously fixed parameter.

        Arguments:
        p -- if supplied, an index or list of indices of parameters to be released. If omitted, release all parameters.
        '''
        
        if p is None:
            self._reset_fcn()
            return
        
        if type(p) == type(list()):
            self._rel_p([p])
        else:
            self._rel_p(p)

    def _rel_p(self, p):
        '''
        Does the recompiling necessary for releasing a parameter.
        '''
        for index in p:
            self.fixedp.remove(index)

        self._fix_p(self.fixedp)
    
    def chi(self, p, xarr, y):
        '''
        The function to be minimized by scipy.optimize algorithms.

        Arguments:

        p    -- a list of parameter values for this iteration
        xarr -- an array of x-coordinates of the data to be fit
        y    -- an array of y-coordinates of the data to be fit
        '''
        
        return [self._fcn(p, x) for x in xarr] - y
        
    def _set_frange_none(self):
        return min(self.x), max(self.x)
                        
    def _set_initpars(self, frange):
        return [1 for i in xrange(self.npar)]

    def fit(self, initpars=None, frange=None, par_names=None):
        '''
        Perform a fit.

        Keyword arguments:

        initpars  -- A list of initial parameters for this fit
        frange    -- The x-range over which the fit should be performed
        par_names -- A list of names for the parameters of this fit
        '''
        
        if par_names is None:
            par_names = self.par_names
        elif len(par_names) != self.npar:
            raise FitException("Length of supplied parameter name list in Fitter.fit() not of expected length ({0} supplied, {1} expected)".format(len(par_names), self.npar)) 

        if initpars is not None and len(initpars) != self.npar:
            raise FitException("Length of supplied initial parameter list in Fitter.fit() not of expected length ({0} supplied, {1} expected)".format(len(initpars), self.npar))

        if frange is None:
            frange = self._set_frange_none()

        if initpars is None:
            initpars = self._set_initpars(frange)


        initpars        = numpy.array(initpars)

        # Epic line of code right here.

        x,y             = map(numpy.array, zip(*[(self.x[i], self.y[i]) for i in xrange(len(self.x)) if self.x[i] >= frange[0] and self.x[i] <= frange[1]]))

        ndof = len(x) - len(initpars)

        init_chi2_dof    = sum([pt**2 for pt in self.chi(initpars, x, y)]) / ndof
        
        fitpars, cov_x, fitinfo, mesg, success = optimize.leastsq(self.chi, initpars[:], args=(x, y), full_output=1, diag=self.weights)

        final_chi2_dof   = sum([pt**2 for pt in self.chi(fitpars, x, y)]) / ndof
        
        self.stats = FitStatistics(frange, x, y, initpars, fitpars, init_chi2_dof,
                                   final_chi2_dof, ndof, success, mesg, cov_x, par_names,)
        return success


        

    def draw(self, name='MyFit', ext='png', plotarg='r-', plotlw=2,  pars=None, same=False, data=True, fit=True, stats=True, hold=False,
             **kwargs):
        '''
        Draw the data and the fitted function.

        Keyword arguments:

        name        -- a name for the fit. The output file will have the form ``name.ext``.
        ext         -- the desired image extension. Allowed values are ``png``, ``gif``, and ``jpg``.
        plotarg     -- a Matplotlib-like string defining how the fit function should be drawn.
        plotlw      -- the width of the line of the drawn function
        pars        -- the parameters with which the function should be drawn (defaults to the last fit parameters)
        same        -- draw without clearing the canvas
        data        -- draw the data points
        fit         -- draw the resulting fit
        stats       -- write the fit parameters and statistics to the image
        hold        -- do not automatically write to file

        '''
        orig_attrs  = dict()
        local_attrs = list()
        for key in kwargs:
            print "BLARG!", key, kwargs[key]
            try:
                orig_attrs[key] = getattr(self, key)
            except AttributeError:
                local_attrs.append(key)

            setattr(self, key, kwargs[key])
        
        for key in kwargs:
            print "BEEP!", getattr(self, key)

        if pars is None:
            pars = self.stats.fitpars
            
        if not same:
            pyplot.clf()
            if stats:
                pyplot.subplots_adjust(right=0.6)

        if data:
            if sum(self.weights) == len(self.weights):
                pyplot.plot(self.x, self.y, 'ko')
                pyplot.plot(self.stats.x, self.stats.y, 'ro')

            else:
                pyplot.errorbar(self.x, self.y, yerr=1.0/self.weights, xerr=None, ecolor='k', mfc='k', mec='r', fmt='o')
                pyplot.errorbar(self.x, self.y, yerr=1.0/self.weights, xerr=None, ecolor='r', mfc='r', mec='k', fmt='o')
                
        if fit:
            if len(self.x) < 50:
                x = numpy.linspace(self.stats.frange[0], self.stats.frange[1])
            else:
                x = self.x

            pyplot.plot(x, [self._fcn(pars, xpt) for xpt in x],
                        plotarg,
                        label='fit',
                        linewidth=plotlw)
        
        if stats:
            self._draw_stats(pars)

        pyplot.title(self.title)
        pyplot.xlabel(self.xlabel)
        pyplot.ylabel(self.ylabel)

        if not hold:
            return self.save_and_clear(name=name)

    def save_and_clear(self, **kwargs):
        '''
        Save the contents of the canvas to file and then clear the canvas.

        Keyword arguments:

        name -- name of the output file. Valid extensions are ``png``, ``gif``, or ``jpg``.
        '''
        orig_attrs  = dict()
        local_attrs = list()
        for key in kwargs:
            try:
                orig_attrs[key] = getattr(self, key)
            except AttributeError:
                local_attrs.append(key)

            setattr(self, key, kwargs[key])
        if self.name == None:
            self.name = 'MyFit'

        retname = None
        if self.name[-3:].lower() in 'png jpg gif'.split():
            pyplot.savefig(self.name)
            retname = self.name
        else:
            outname = '.'.join([self.name, self.ext])
            pyplot.savefig(outname)
            retname = outname

        for key in orig_attrs:
            setattr(self, key, orig_attrs[key])
        
        for attr in local_attrs:
            delattr(self, attr)

        pyplot.clf()
        pyplot.subplots_adjust(right=0.9)

        return retname

    def _draw_stats(self, pars):
        '''
        Draw the fit parameters and statistics on the current canvas.

        Arguments:

        pars -- The list of parameters to be written out to the canvas.
        '''
        
        fitstr = ""

        if self.stats.par_names is not None:
            for i in xrange(len(pars)):
                fitstr += "  {0} = {1:.3e}\n".format(self.stats.par_names[i], pars[i])
        else:
            for i in xrange(len(pars)):
                fitstr += "  $p_{0}$ = {1:.3e}\n".format(i, pars[i])

        pars_str = "Parameters:\n"
        pars_str += "  " + self.fcn_latex + '\n\n'
        pars_str += fitstr

        stats_str = "# of points: {0:1d}\n".format(len(self.x))
        stats_str += "$(\\chi^2/d.o.f.)_i$: {0:.1e}\n".format(self.stats.init_chi2_dof)
        stats_str += "$(\\chi^2/d.o.f.)_f$: {0:.1e}\n".format(self.stats.final_chi2_dof)
        stats_str += "n.d.o.f.: " + str(self.stats.ndof) + '\n'

        pyplot.figtext(0.65, self._par_text_offset, pars_str, size='x-large')
        pyplot.figtext(0.65, 0.045, stats_str, size='x-large')

